# -*- coding: utf-8 -*-
from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import os


# Transformer模型定义
class TransformerModel(nn.Module):
    def __init__(self, model_dim, n_heads, n_layers, output_dim):
        super(TransformerModel, self).__init__()
        self.transformer_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(model_dim, n_heads), num_layers=n_layers
        )
        # 修改分类层为多层全连接层
        self.fc1 = nn.Linear(model_dim, 328)  # 第一个全连接层
        self.fc2 = nn.Linear(328, 64)  # 第二个全连接层
        self.fc3 = nn.Linear(64, output_dim)  # 输出层

    def forward(self, x):
        x = self.transformer_encoder(x)  # 直接传入矩阵，经过Transformer编码器
        x = x.mean(dim=1)  # 对输出取平均，得到句子的表示
        x = F.relu(self.fc1(x))  # 第一个全连接层，使用ReLU激活
        x = F.relu(self.fc2(x))  # 第二个全连接层，使用ReLU激活
        x = self.fc3(x)  # 输出层
        return x


# 模型评估函数
def evaluate_model(model, test_loader, epoch, device):
    model.eval()
    total_correct = 0
    total_samples = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, dim=1)  # 选择每个样本的最大类别
            labels = torch.argmax(labels, dim=1)
            total_correct += (predicted == labels).sum().item()
            total_samples += labels.size(0)

    accuracy = total_correct / total_samples
    print(f'Epoch [{epoch+1}] Accuracy: {accuracy:.4f}')



# 模型参数
num_epochs = 500
model_dim = 300  # 输入矩阵的维度
n_heads = 2  # 注意力头数
n_layers = 2  # Transformer层数
output_dim = 2  # 二分类
batch_size = 16  # 批量大小

# 加载数据
data = np.load('../data/data_matrix.npz')
sentences_array = data['sentences']  # 句子矩阵
tags_array = data['tags']

# 过滤掉空句子
filtered_sentences = [sentence for sentence in sentences_array if len(sentence) > 0]
filtered_tags = [tags_array[i] for i in range(len(sentences_array)) if len(sentences_array[i]) > 0]

# 确保输入数据是统一长度（这里需要做填充）
max_len = max(len(sentence) for sentence in filtered_sentences)
padded_sentences = np.array(
    [np.pad(sentence, ((0, max_len - len(sentence)), (0, 0)), 'constant') for sentence in filtered_sentences])

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(padded_sentences, filtered_tags, test_size=0.2, random_state=42)

# 创建DataLoader
train_dataset = TensorDataset(torch.tensor(X_train, dtype=torch.float32), torch.tensor(y_train, dtype=torch.float32))
test_dataset = TensorDataset(torch.tensor(X_test, dtype=torch.float32), torch.tensor(y_test, dtype=torch.float32))

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 创建模型
model = TransformerModel(model_dim=model_dim, n_heads=n_heads, n_layers=n_layers, output_dim=output_dim)
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)  # 将模型移动到GPU

# 训练循环
for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0  # 记录每个epoch的总损失
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        epoch_loss += loss.item()

    if (epoch + 1) % 10 == 0:
        torch.save(model.state_dict(), 'transformer_model.pth')

    # print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss / len(train_loader):.4f}')

    # 评估模型
    evaluate_model(model, test_loader, epoch, device)

# 训练循环结束后，保存模型
torch.save(model.state_dict(), 'transformer_model.pth')
print("模型已保存为 transformer_model.pth")
